#
#  Copyright 2021 The FATE Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import errno
import os
import subprocess
import time

import psutil

from pyoi.util import base_utils
from pyoi.util.log_utils import getLogger

logger = getLogger()


def run_subprocess(process_cmd, added_env: dict = None, config_dir=None, log_dir=None, process_name="", process_id=""):
    process_cmd = [str(cmd) for cmd in process_cmd]
    logger.info("start process command: \n{}".format(" ".join(process_cmd)))

    os.makedirs(config_dir, exist_ok=True)
    if not log_dir:
        log_dir = config_dir
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    if not process_id:
        process_id = base_utils.current_timestamp()
    std_path = get_std_path(log_dir=log_dir, process_name=process_name, process_id=process_id)
    std = open(std_path, 'w')
    pid_path = os.path.join(config_dir, f"{process_name}_pid")

    if os.name == 'nt':
        startupinfo = subprocess.STARTUPINFO()
        startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        startupinfo.wShowWindow = subprocess.SW_HIDE
    else:
        startupinfo = None
    subprocess_env = os.environ.copy()
    if added_env:
        for name, value in added_env.items():
            if name.endswith("PATH"):
                subprocess_env[name] = subprocess_env.get(name, "") + f":{value}"
            else:
                subprocess_env[name] = value
    if "CLASSPATH" in subprocess_env:
        del subprocess_env["CLASSPATH"]
    p = subprocess.Popen(process_cmd,
                         stdout=std,
                         stderr=std,
                         startupinfo=startupinfo,
                         cwd=config_dir,
                         env=subprocess_env
                         )
    with open(pid_path, 'w') as f:
        f.truncate()
        f.write(str(p.pid) + "\n")
        f.flush()
    logger.info(f"start process successfully, pid: {p.pid}, std log path: {std_path}")
    return p


def check_process(process: subprocess.Popen):
    return process.poll() is None


def kill_process(process: subprocess.Popen, expected_cmdline: list):
    try:
        pid = process.pid
        process.kill()
        time.sleep(0.01)
        if check_process_by_pid(pid, expected_cmdline=expected_cmdline):
            kill_process_by_pid(pid=pid, expected_cmdline=expected_cmdline)
    except Exception as e:
        logger.warning(e)


def check_process_by_pid(pid, expected_cmdline: list = None):
    if pid < 0:
        return False
    if pid == 0:
        raise ValueError('invalid PID 0')
    try:
        os.kill(pid, 0)
    except OSError as err:
        if err.errno == errno.ESRCH:
            # ESRCH == No such process
            ret = False
        elif err.errno == errno.EPERM:
            # EPERM clearly means there's a process to deny access to
            ret = True
        else:
            # According to "man 2 kill" possible error values are
            # (EINVAL, EPERM, ESRCH)
            raise
    else:
        ret = True
    if ret and expected_cmdline is not None:
        p = get_process_instance(pid)
        try:
            return check_process_by_cmdline(actual=p.cmdline(), expected=expected_cmdline)
        except psutil.ZombieProcess as e:
            logger.warning(f"process {pid} zombie")
            return False
    else:
        return ret


def check_process_by_cmdline(actual: list, expected: list):
    if len(actual) != len(expected):
        return False
    for i, v in enumerate(actual):
        if str(v) != str(expected[i]):
            return False
    else:
        return True


def get_std_path(log_dir, process_name="", process_id=""):
    std_log_path = f"{process_name}_{process_id}_std.log"
    return os.path.join(log_dir, std_log_path)


def get_subprocess_std(log_dir, process_name="", process_id=""):
    with open(get_std_path(log_dir, process_name, process_id), "r") as fr:
        text = fr.read()
    return text


def wait_child_process(signum, frame):
    try:
        while True:
            child_pid, status = os.waitpid(-1, os.WNOHANG)
            if child_pid == 0:
                logger.info('no child process was immediately available')
                break
            exitcode = status >> 8
            logger.info(f'child process {child_pid} exit with exitcode {exitcode}')
    except OSError as e:
        if e.errno == errno.ECHILD:
            logger.info('current process has no existing unwaited-for child processes.')
        else:
            raise


def kill_process_by_pid(process: psutil.Process = None, pid: int = None, expected_cmdline: list = None):
    process = process if process is not None else get_process_instance(pid)
    if process is None:
        return
    for child in process.children(recursive=True):
        try:
            if check_process_by_pid(pid=child.pid):
                child.kill()
        except psutil.NoSuchProcess:
            pass
        except Exception as e:
            logger.warning(f"kill {child.pid} process failed: {e}", exc_info=True)
    if check_process_by_pid(pid=process.pid, expected_cmdline=expected_cmdline):
        process.kill()


def get_process_instance(pid: int):
    try:
        return psutil.Process(int(pid))
    except psutil.NoSuchProcess:
        logger.warning(f"no such process {pid}")
        return
    except Exception as e:
        raise e
